import pandas as pd
from yi_util import *
from yi_util import divide_by_100,get_last_days_df
from data_cal import cal_gt4

'''============================================ 逻辑处理 BEGIN ============================================'''
'''筛选股票列表逻辑'''
def filter_s_list(code_df):
    code_df.rename(columns={'symbol':'code','display_name':'name'},inplace=True)
    filtered_df = add_zt_type(code_df)
    return filtered_df


    
def filter_days(day_df):
    day_df['turnover'] = day_df['turnover']/100000000 #除以1亿，方便后期查看
    day_df= day_df.rename(columns={'turnover_rate':'t_rate'})
    return day_df





def filter_flow(flow_df):
    # 净特大单 = 主被动买入特大单 - 主被动卖出特大单
    flow_df['xl'] = (flow_df['act_buy_xl'] + flow_df['pas_buy_xl']) - (flow_df['act_sell_xl'] + flow_df['pas_sell_xl'])
    flow_df = flow_df.applymap(divide_by_100)
    flow_df = flow_df[['date', 'code', 'xl']]
    flow_df = flow_df.sort_values(by=['date'], ascending=True)
    flow_df = flow_df.round(2).reset_index(drop=True)
    flow_df = flow_df.groupby('code').apply(calculate_yoy_growth).reset_index(drop=True).round(2)
    return flow_df


def filter_factor(factor_df):
    factor_df=factor_df.rename(columns={'factor_symbol':'code','factor_date':'date','factor_current_market_cap':'m_c',})
    # factor_df = factor_df[['date','code','m_c']]
    factor_df['m_c']=factor_df['m_c']/1000000#（100万）
    return factor_df

''' 流通市值为空时，自己计算数据 '''
def filter_factor_empty(day_df,factor_df):
    last_factor_df = get_last_day_df(factor_df,2)
    last_day_df = get_last_day_df(day_df,1)[['date','code','quote_rate']]
    day_factore_df =  pd.merge(last_factor_df[['code','m_c']], last_day_df, on=['code'])
    day_factore_df['m_c'] = day_factore_df['m_c']*(1 + day_factore_df['quote_rate']/100).round(2)
    return day_factore_df

''' 统计某些列近几日累加值 '''
def filter_xl_sum(df, window_size=4):
    df['date'] = pd.to_datetime(df['date'])
    df = df.sort_values(['code', 'date'])
    df['xlt_sum'] = df.groupby('code')['xl_turnover'].rolling(window=window_size).sum().droplevel(0)
    df['xlmc_sum'] = df.groupby('code')['xl_m_c'].rolling(window=window_size).sum().droplevel(0)
    df = df.round(2)
    df = get_last_days_df(df,20)
    return df.reset_index(drop=True)

def filter_final_merge(df):
    min_mc = 800 # 8~10亿（单位是百万）
    max_mc = 50000 # 200亿（单位是百万）
    df = df[(df['m_c'] >= min_mc)]
    df = df[(df['m_c'] <= max_mc)]

    # 定义基础列列表（不包含gn_status）
    base_columns = [
        'date', 'code', 'quote_rate', 'xl', 'xl_turnover',
        # 'xlt_sum',
        # 'xlmc_sum',
        'd_week', 'd_ma15', 'r2', 'r5', 'm_raise', 'zt_num', 
        'gt4', 'gt9', 'd_ma3_5', 'd_ma3', 'd_ma5'
    ]
    
    # 如果存在gn_status则追加到列列表
    if 'gn_status' in df.columns:
        base_columns.append('gn_status')
    
    # 按最终列列表选择数据
    df = df[base_columns]
    return df
''' 概念涨停率筛选 '''
def filter_gn_ztratio(zt_ratio_df):
     # ==合格阈值==
    zt_ratio_df = zt_ratio_df[(zt_ratio_df['pass_num'] >= 5)]
    zt_ratio_df = zt_ratio_df[(zt_ratio_df['zt_ratio'] >= 10)]
    zt_ratio_df = zt_ratio_df[(zt_ratio_df['item_count'] <= 200)]
    return zt_ratio_df

''' 单只个股筛选 '''
def filter_single_code(code,df):
    code = append_suffix(code)
    df = df[df['code'] ==code]
    df = r_index(df)
    return df

''' 过滤 day数据： 大于15日均线；价格：3~40，过去10日gt4出现过至少一次'''
def filter_day_raise_ma_close(df):
    last_df = get_last_day_df(df,1)
    last_df = last_df[(last_df['gt4'])]
    last_df = last_df[(last_df['d_ma15'] >= 1)]
    last_df = last_df[(last_df['close'] < 40)]
    last_df = last_df[(last_df['close'] > 3)]
    df = df[df['code'].isin(last_df['code'])]
    return df


def mark_consecutive_zt(df: pd.DataFrame) -> pd.DataFrame:
    """
    标记连续涨停天数
    参数：
        df : 必须包含列 ['date', 'code', 'is_zt']
    返回：
        新增 zt_count 列的 DataFrame，取值规则：
        - 1: 首日涨停
        - 2: 连续两日涨停
        - 3: 连续三日及以上涨停
    """
    # 创建副本避免修改原数据
    df = df.copy()
    
    # 按股票代码和时间排序（处理乱序数据）
    df = df.sort_values(by=['code', 'date']).reset_index(drop=True)
    
    # 计算连续涨停次数
    def _count_consecutive(group):
        # 生成连续涨停标记组
        group['cum_zt'] = (group['is_zt'].astype(int)
                          .groupby((group['is_zt'] != group['is_zt'].shift()).cumsum())
                          .cumcount() + 1)
        # 非涨停日重置为0
        group['cum_zt'] = group['cum_zt'].where(group['is_zt'], 0)
        return group
    
    # 分组处理每个股票
    df = df.groupby('code', group_keys=False).apply(_count_consecutive)
    
    # 映射到分类值
    df['zt_count'] = df['cum_zt'].apply(
        lambda x: min(3, x) if x >= 1 else 0  # 连续3天及以上统一标记为3
    )
    
    # 清理中间列
    df.drop(columns=['cum_zt'], inplace=True, errors='ignore')
    
    return df

def add_zt_type(df: pd.DataFrame, 
                code_col: str = 'code', 
                name_col: str = 'name', 
                result_col: str = 'zt_type') -> pd.DataFrame:
    """
    批量添加股票涨停类型标记
    
    Parameters:
        df: 原始DataFrame
        code_col: 股票代码列名（默认'code'）
        name_col: 股票名称列名（默认'name'）
        result_col: 结果列名（默认'zt_type'）
    
    Returns:
        带zt_type标记的DataFrame
    """
    # 类型转换确保处理字符串
    df = df.copy()
    df[code_col] = df[code_col].astype(str)
    
    # 条件优先级顺序：ST标记 > 代码前缀
    conditions = [
        # ST标记检测（网页5的布尔索引思路）
        df[name_col].str.contains('ST', case=False),  
        # 代码前缀检测（网页7的startswith用法）
        df[code_col].str.startswith(('60', '00')),     # 10cm
        df[code_col].str.startswith(('30', '68')),     # 20cm 
        df[code_col].str.startswith(('43', '83', '92')) # 30cm
    ]
    choices = [5, 10, 20, 30]
    
    # 批量赋值（网页3的np.select应用）
    df[result_col] = np.select(conditions, choices, default=np.nan).astype('int64')
    return df


'''============================================ 逻辑处理 END ============================================'''